import torch
from torchaudio.models import Tacotron2
import torchaudio
from torchaudio.models.tacotron2 import Tacotron2

# 修改后的 Tacotron2 初始化代码
model = Tacotron2(
    n_symbol=100,  # 符号数量
    n_mels=80,  # 梅尔频谱通道数
    symbol_embedding_dim=512,  # 符号嵌入维度
    encoder_embedding_dim=512,  # 编码器嵌入维度
    encoder_n_convolution=3,  # 编码器卷积层数
    encoder_kernel_size=5,  # 编码器卷积核大小
    decoder_rnn_dim=1024,  # 解码器RNN维度
    decoder_max_step=2000,  # 最大解码步数
    decoder_dropout=0.1,  # 解码器dropout概率
    decoder_early_stopping=True,  # 是否提前停止解码
    attention_rnn_dim=1024,  # 注意力RNN维度
    attention_hidden_dim=128,  # 注意力隐藏层维度
    attention_location_n_filter=32,  # 注意力位置滤波器数量
    attention_location_kernel_size=31,  # 注意力位置卷积核大小
    attention_dropout=0.1,  # 注意力dropout概率
    prenet_dim=256,  # 预网络维度
    postnet_n_convolution=5,  # 后处理网络卷积层数
    postnet_kernel_size=5,  # 后处理网络卷积核大小
    postnet_embedding_dim=512,  # 后处理网络嵌入维度
    gate_threshold=0.5  # 门限阈值
)

# 加载预训练权重
checkpoint = torch.load("pretrained_ljs.pth")
checkpoint = torch.load("pretrained_ljs.pth")
if 'state_dict' in checkpoint:
    model.load_state_dict(checkpoint['state_dict'])
else:
    # 如果 checkpoint 直接包含模型参数
    model.load_state_dict(checkpoint)

# 定义文本预处理函数
def preprocess_text(text):
    # 将文本转换为模型可接受的输入格式
    # 这里假设需要将文本转换为索引序列
    # 你需要根据实际模型需求实现具体的预处理逻辑
    return [ord(char) for char in text]  # 简单示例：将字符转换为ASCII码

# 定义流式处理的块大小
chunk_size = 10  # 每次处理10个字符

# ... existing code ...

def text_to_speech(text, streaming=False):
    # 文本预处理
    sequence = preprocess_text(text)

    sequence_tensor = torch.tensor(sequence, dtype=torch.long)  # 将列表转换为Tensor

    # 流式输出
    if streaming:
        for i in range(0, len(sequence_tensor), chunk_size):
            chunk = sequence_tensor[i:i + chunk_size]
            mel = model.infer(chunk)
            yield mel
    else:
        mel = model.infer(sequence_tensor)
        return mel


def test_text_to_speech():
    # 测试文本
    test_text = "你好，这是一个测试文本。"

    # 测试非流式模式
    print("测试非流式模式...")
    mel_spectrogram = text_to_speech(test_text, streaming=False)
    print(f"返回类型: {type(mel_spectrogram)}")
    print(f"梅尔频谱图形状: {mel_spectrogram.shape}")

    # 保存生成的音频
    torchaudio.save("output_non_streaming.wav", mel_spectrogram, sample_rate=24000)

    # 测试流式模式
    print("\n测试流式模式...")
    stream = text_to_speech(test_text, streaming=True)
    for i, chunk in enumerate(stream):
        print(f"第 {i + 1} 块, 类型: {type(chunk)}, 形状: {chunk.shape}")
        # 保存每个音频块
        torchaudio.save(f"output_streaming_chunk_{i + 1}.wav", chunk, sample_rate=24000)


if __name__ == "__main__":
    test_text_to_speech()
